import torch
import matplotlib.pyplot as plt

"""
NAG算法：考虑累积的动量 和 预测点的梯度，更新动量，再去更新参数。
相比直接Momentum动量法，能“预见”未来，更好的收敛和减少震荡。
"""

# 1. 定义损失函数
def loss_fn(w1, w2):
    return w1 ** 2 + 2 * w2 ** 2

# 2. 超参数
lr = 0.05
Epochs = 20
beta = 0.9
w1 = -1
w2 = 1
v1 = 0
v2 = 0

''' 绘制等高线图 '''
x1 = torch.linspace(-1, 1, 100)
x2 = torch.linspace(-1, 1, 100)
xx1, xx2 = torch.meshgrid(x1, x2, indexing="ij")
loss = loss_fn(xx1, xx2)
fig = plt.figure("NAG")
ax = fig.add_subplot()
ax.contour(xx1, xx2, loss)
# 定义一个列表，用于存储梯度下降的路径点
points = []

# 3. 循环训练
for epoch in range(Epochs):
    points.append([w1, w2])     # 保存参数点
    loss = loss_fn(w1, w2)
    print(loss)
    # ==== NAG算法 ====
    # 根据累积的动量信息，预测下一个可能的参数位置
    w1_pre = w1 + beta * v1
    w2_pre = w2 + beta * v2
    # 计算 【预测到的下一个可能的参数位置】 的梯度
    g1_pre = 2 * w1_pre
    g2_pre = 4 * w2_pre
    # 更新动量
    v1 = beta * v1 - lr * g1_pre
    v2 = beta * v2 - lr * g2_pre
    # 更新参数
    w1 += v1
    w2 += v2

points = torch.tensor(points)
ax.plot(points[:,0], points[:,1], "ko-")
plt.show()
